﻿using System;
using System.Runtime.Intrinsics;
using System.Runtime.Intrinsics.X86;
using System.Runtime.Intrinsics.Arm;


namespace VectorCalculator
{

    /// <summary>
    /// 跨平台SIMD向量计算帮助类
    /// 支持x64 (AVX2/SSE4.1) 和 ARM64 (NEON) 指令集，自动回退到标量运算
    /// </summary>
    public static class CalcWithSIMD
    {
        private const int AVX_VECTOR_SIZE = 8;   // AVX2处理8个float
        private const int NEON_VECTOR_SIZE = 4;  // NEON处理4个float
        private const int SSE_VECTOR_SIZE = 4;   // SSE处理4个float

        #region 哈达玛积 (逐元素相乘)

        /// <summary>
        /// 向量哈达玛积（逐元素相乘）
        /// </summary>
        public static float[] HadamardProduct(float[] vector1, float[] vector2)
        {
            if (vector1 == null || vector2 == null)
                throw new ArgumentException("向量不能为空");
            if (vector1.Length != vector2.Length)
                throw new ArgumentException("向量长度必须一致");

            float[] result = new float[vector1.Length];

            if (Avx.IsSupported && vector1.Length >= AVX_VECTOR_SIZE)
            {
                HadamardProductAvx(vector1, vector2, result);
            }
            else if (AdvSimd.IsSupported && vector1.Length >= NEON_VECTOR_SIZE)
            {
                HadamardProductNeon(vector1, vector2, result);
            }
            else if (Sse.IsSupported && vector1.Length >= SSE_VECTOR_SIZE)
            {
                HadamardProductSse(vector1, vector2, result);
            }
            else
            {
                HadamardProductScalar(vector1, vector2, result);
            }

            return result;
        }

        private static unsafe void HadamardProductAvx(float[] vector1, float[] vector2, float[] result)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % AVX_VECTOR_SIZE);

            fixed (float* pV1 = vector1, pV2 = vector2, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += AVX_VECTOR_SIZE)
                {
                    Vector256<float> v1 = Avx.LoadVector256(pV1 + i);
                    Vector256<float> v2 = Avx.LoadVector256(pV2 + i);
                    Vector256<float> product = Avx.Multiply(v1, v2);
                    Avx.Store(pResult + i, product);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector1[i] * vector2[i];
            }
        }

        private static unsafe void HadamardProductNeon(float[] vector1, float[] vector2, float[] result)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % NEON_VECTOR_SIZE);

            fixed (float* pV1 = vector1, pV2 = vector2, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += NEON_VECTOR_SIZE)
                {
                    Vector128<float> v1 = AdvSimd.LoadVector128(pV1 + i);
                    Vector128<float> v2 = AdvSimd.LoadVector128(pV2 + i);
                    Vector128<float> product = AdvSimd.Multiply(v1, v2);
                    AdvSimd.Store(pResult + i, product);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector1[i] * vector2[i];
            }
        }

        private static unsafe void HadamardProductSse(float[] vector1, float[] vector2, float[] result)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % SSE_VECTOR_SIZE);

            fixed (float* pV1 = vector1, pV2 = vector2, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += SSE_VECTOR_SIZE)
                {
                    Vector128<float> v1 = Sse.LoadVector128(pV1 + i);
                    Vector128<float> v2 = Sse.LoadVector128(pV2 + i);
                    Vector128<float> product = Sse.Multiply(v1, v2);
                    Sse.Store(pResult + i, product);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector1[i] * vector2[i];
            }
        }

        private static void HadamardProductScalar(float[] vector1, float[] vector2, float[] result)
        {
            for (int i = 0; i < vector1.Length; i++)
            {
                result[i] = vector1[i] * vector2[i];
            }
        }

        #endregion

        #region 向量缩放 (标量乘积)

        /// <summary>
        /// 向量缩放（标量乘积）
        /// </summary>
        public static float[] Scale(float[] vector, float scalar)
        {
            if (vector == null)
                throw new ArgumentException("向量不能为空");

            float[] result = new float[vector.Length];

            if (Avx.IsSupported && vector.Length >= AVX_VECTOR_SIZE)
            {
                ScaleAvx(vector, scalar, result);
            }
            else if (AdvSimd.IsSupported && vector.Length >= NEON_VECTOR_SIZE)
            {
                ScaleNeon(vector, scalar, result);
            }
            else if (Sse.IsSupported && vector.Length >= SSE_VECTOR_SIZE)
            {
                ScaleSse(vector, scalar, result);
            }
            else
            {
                ScaleScalar(vector, scalar, result);
            }

            return result;
        }

        private static unsafe void ScaleAvx(float[] vector, float scalar, float[] result)
        {
            int length = vector.Length;
            int vectorizedLength = length - (length % AVX_VECTOR_SIZE);
            Vector256<float> scalarVec = Vector256.Create(scalar);

            fixed (float* pVector = vector, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += AVX_VECTOR_SIZE)
                {
                    Vector256<float> v = Avx.LoadVector256(pVector + i);
                    Vector256<float> scaled = Avx.Multiply(v, scalarVec);
                    Avx.Store(pResult + i, scaled);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector[i] * scalar;
            }
        }

        private static unsafe void ScaleNeon(float[] vector, float scalar, float[] result)
        {
            int length = vector.Length;
            int vectorizedLength = length - (length % NEON_VECTOR_SIZE);
            Vector128<float> scalarVec = Vector128.Create(scalar);

            fixed (float* pVector = vector, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += NEON_VECTOR_SIZE)
                {
                    Vector128<float> v = AdvSimd.LoadVector128(pVector + i);
                    Vector128<float> scaled = AdvSimd.Multiply(v, scalarVec);
                    AdvSimd.Store(pResult + i, scaled);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector[i] * scalar;
            }
        }

        private static unsafe void ScaleSse(float[] vector, float scalar, float[] result)
        {
            int length = vector.Length;
            int vectorizedLength = length - (length % SSE_VECTOR_SIZE);
            Vector128<float> scalarVec = Vector128.Create(scalar);

            fixed (float* pVector = vector, pResult = result)
            {
                for (int i = 0; i < vectorizedLength; i += SSE_VECTOR_SIZE)
                {
                    Vector128<float> v = Sse.LoadVector128(pVector + i);
                    Vector128<float> scaled = Sse.Multiply(v, scalarVec);
                    Sse.Store(pResult + i, scaled);
                }
            }

            for (int i = vectorizedLength; i < length; i++)
            {
                result[i] = vector[i] * scalar;
            }
        }

        private static void ScaleScalar(float[] vector, float scalar, float[] result)
        {
            for (int i = 0; i < vector.Length; i++)
            {
                result[i] = vector[i] * scalar;
            }
        }

        #endregion

        #region L2距离

        /// <summary>
        /// 计算两个向量的L2距离（欧几里德距离）
        /// </summary>
        public static float L2Distance(float[] vector1, float[] vector2)
        {
            if (vector1 == null || vector2 == null)
                throw new ArgumentException("向量不能为空");
            if (vector1.Length != vector2.Length)
                throw new ArgumentException("向量长度必须一致");

            if (Avx.IsSupported && vector1.Length >= AVX_VECTOR_SIZE)
            {
                return L2DistanceAvx(vector1, vector2);
            }
            else if (AdvSimd.IsSupported && vector1.Length >= NEON_VECTOR_SIZE)
            {
                return L2DistanceNeon(vector1, vector2);
            }
            else if (Sse.IsSupported && vector1.Length >= SSE_VECTOR_SIZE)
            {
                return L2DistanceSse(vector1, vector2);
            }
            else
            {
                return L2DistanceScalar(vector1, vector2);
            }
        }

        private static unsafe float L2DistanceAvx(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % AVX_VECTOR_SIZE);
            Vector256<float> sumOfSquares = Vector256<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += AVX_VECTOR_SIZE)
                {
                    Vector256<float> v1 = Avx.LoadVector256(pV1 + i);
                    Vector256<float> v2 = Avx.LoadVector256(pV2 + i);
                    Vector256<float> diff = Avx.Subtract(v1, v2);
                    sumOfSquares = Avx.Add(sumOfSquares, Avx.Multiply(diff, diff));
                }
            }

            float totalSumOfSquares = HorizontalSumAvx(sumOfSquares);

            for (int i = vectorizedLength; i < length; i++)
            {
                float diff = vector1[i] - vector2[i];
                totalSumOfSquares += diff * diff;
            }

            return (float)Math.Sqrt(totalSumOfSquares);
        }

        private static unsafe float L2DistanceNeon(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % NEON_VECTOR_SIZE);
            Vector128<float> sumOfSquares = Vector128<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += NEON_VECTOR_SIZE)
                {
                    Vector128<float> v1 = AdvSimd.LoadVector128(pV1 + i);
                    Vector128<float> v2 = AdvSimd.LoadVector128(pV2 + i);
                    Vector128<float> diff = AdvSimd.Subtract(v1, v2);
                    sumOfSquares = AdvSimd.Add(sumOfSquares, AdvSimd.Multiply(diff, diff));
                }
            }

            float totalSumOfSquares = HorizontalSumNeon(sumOfSquares);

            for (int i = vectorizedLength; i < length; i++)
            {
                float diff = vector1[i] - vector2[i];
                totalSumOfSquares += diff * diff;
            }

            return (float)Math.Sqrt(totalSumOfSquares);
        }

        private static unsafe float L2DistanceSse(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % SSE_VECTOR_SIZE);
            Vector128<float> sumOfSquares = Vector128<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += SSE_VECTOR_SIZE)
                {
                    Vector128<float> v1 = Sse.LoadVector128(pV1 + i);
                    Vector128<float> v2 = Sse.LoadVector128(pV2 + i);
                    Vector128<float> diff = Sse.Subtract(v1, v2);
                    sumOfSquares = Sse.Add(sumOfSquares, Sse.Multiply(diff, diff));
                }
            }

            float totalSumOfSquares = HorizontalSumSse(sumOfSquares);

            for (int i = vectorizedLength; i < length; i++)
            {
                float diff = vector1[i] - vector2[i];
                totalSumOfSquares += diff * diff;
            }

            return (float)Math.Sqrt(totalSumOfSquares);
        }

        private static float L2DistanceScalar(float[] vector1, float[] vector2)
        {
            float sumOfSquares = 0f;
            for (int i = 0; i < vector1.Length; i++)
            {
                float diff = vector1[i] - vector2[i];
                sumOfSquares += diff * diff;
            }
            return (float)Math.Sqrt(sumOfSquares);
        }

        #endregion

        #region 余弦相似度

        /// <summary>
        /// 计算两个向量的余弦相似度
        /// </summary>
        public static float CosineSimilarity(float[] vector1, float[] vector2)
        {
            if (vector1 == null || vector2 == null)
                throw new ArgumentException("向量不能为空");
            if (vector1.Length == 0 || vector2.Length == 0)
                throw new ArgumentException("向量长度不能为0");
            if (vector1.Length != vector2.Length)
                throw new ArgumentException("向量长度必须一致");

            if (Avx.IsSupported && vector1.Length >= AVX_VECTOR_SIZE)
            {
                return CosineSimilarityAvx(vector1, vector2);
            }
            else if (AdvSimd.IsSupported && vector1.Length >= NEON_VECTOR_SIZE)
            {
                return CosineSimilarityNeon(vector1, vector2);
            }
            else if (Sse.IsSupported && vector1.Length >= SSE_VECTOR_SIZE)
            {
                return CosineSimilaritySse(vector1, vector2);
            }
            else
            {
                return CosineSimilarityScalar(vector1, vector2);
            }
        }

        private static unsafe float CosineSimilarityAvx(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % AVX_VECTOR_SIZE);

            Vector256<float> dotProductVec = Vector256<float>.Zero;
            Vector256<float> normAVec = Vector256<float>.Zero;
            Vector256<float> normBVec = Vector256<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += AVX_VECTOR_SIZE)
                {
                    Vector256<float> v1 = Avx.LoadVector256(pV1 + i);
                    Vector256<float> v2 = Avx.LoadVector256(pV2 + i);

                    dotProductVec = Avx.Add(dotProductVec, Avx.Multiply(v1, v2));
                    normAVec = Avx.Add(normAVec, Avx.Multiply(v1, v1));
                    normBVec = Avx.Add(normBVec, Avx.Multiply(v2, v2));
                }
            }

            float dotProduct = HorizontalSumAvx(dotProductVec);
            float normA = HorizontalSumAvx(normAVec);
            float normB = HorizontalSumAvx(normBVec);

            for (int i = vectorizedLength; i < length; i++)
            {
                dotProduct += vector1[i] * vector2[i];
                normA += vector1[i] * vector1[i];
                normB += vector2[i] * vector2[i];
            }

            return dotProduct / ((float)Math.Sqrt(normA) * (float)Math.Sqrt(normB));
        }

        private static unsafe float CosineSimilarityNeon(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % NEON_VECTOR_SIZE);

            Vector128<float> dotProductVec = Vector128<float>.Zero;
            Vector128<float> normAVec = Vector128<float>.Zero;
            Vector128<float> normBVec = Vector128<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += NEON_VECTOR_SIZE)
                {
                    Vector128<float> v1 = AdvSimd.LoadVector128(pV1 + i);
                    Vector128<float> v2 = AdvSimd.LoadVector128(pV2 + i);

                    dotProductVec = AdvSimd.Add(dotProductVec, AdvSimd.Multiply(v1, v2));
                    normAVec = AdvSimd.Add(normAVec, AdvSimd.Multiply(v1, v1));
                    normBVec = AdvSimd.Add(normBVec, AdvSimd.Multiply(v2, v2));
                }
            }

            float dotProduct = HorizontalSumNeon(dotProductVec);
            float normA = HorizontalSumNeon(normAVec);
            float normB = HorizontalSumNeon(normBVec);

            for (int i = vectorizedLength; i < length; i++)
            {
                dotProduct += vector1[i] * vector2[i];
                normA += vector1[i] * vector1[i];
                normB += vector2[i] * vector2[i];
            }

            return dotProduct / ((float)Math.Sqrt(normA) * (float)Math.Sqrt(normB));
        }

        private static unsafe float CosineSimilaritySse(float[] vector1, float[] vector2)
        {
            int length = vector1.Length;
            int vectorizedLength = length - (length % SSE_VECTOR_SIZE);

            Vector128<float> dotProductVec = Vector128<float>.Zero;
            Vector128<float> normAVec = Vector128<float>.Zero;
            Vector128<float> normBVec = Vector128<float>.Zero;

            fixed (float* pV1 = vector1, pV2 = vector2)
            {
                for (int i = 0; i < vectorizedLength; i += SSE_VECTOR_SIZE)
                {
                    Vector128<float> v1 = Sse.LoadVector128(pV1 + i);
                    Vector128<float> v2 = Sse.LoadVector128(pV2 + i);

                    dotProductVec = Sse.Add(dotProductVec, Sse.Multiply(v1, v2));
                    normAVec = Sse.Add(normAVec, Sse.Multiply(v1, v1));
                    normBVec = Sse.Add(normBVec, Sse.Multiply(v2, v2));
                }
            }

            float dotProduct = HorizontalSumSse(dotProductVec);
            float normA = HorizontalSumSse(normAVec);
            float normB = HorizontalSumSse(normBVec);

            for (int i = vectorizedLength; i < length; i++)
            {
                dotProduct += vector1[i] * vector2[i];
                normA += vector1[i] * vector1[i];
                normB += vector2[i] * vector2[i];
            }

            return dotProduct / ((float)Math.Sqrt(normA) * (float)Math.Sqrt(normB));
        }

        private static float CosineSimilarityScalar(float[] vector1, float[] vector2)
        {
            float dotProduct = 0f;
            float normA = 0f;
            float normB = 0f;

            for (int i = 0; i < vector1.Length; i++)
            {
                dotProduct += vector1[i] * vector2[i];
                normA += vector1[i] * vector1[i];
                normB += vector2[i] * vector2[i];
            }

            return dotProduct / ((float)Math.Sqrt(normA) * (float)Math.Sqrt(normB));
        }

        #endregion

        #region 辅助方法

        /// <summary>
        /// AVX向量水平求和
        /// </summary>
        private static float HorizontalSumAvx(Vector256<float> vector)
        {
            Vector128<float> high = vector.GetUpper();
            Vector128<float> low = vector.GetLower();
            Vector128<float> sum = Sse.Add(high, low);
            return HorizontalSumSse(sum);
        }

        /// <summary>
        /// ARM NEON向量水平求和
        /// </summary>
        private static unsafe float HorizontalSumNeon(Vector128<float> vector)
        {
            // 手动提取各个元素进行求和（最可靠的方法）
            float* ptr = (float*)&vector;
            return ptr[0] + ptr[1] + ptr[2] + ptr[3];
        }

        /// <summary>
        /// SSE向量水平求和
        /// </summary>
        private static float HorizontalSumSse(Vector128<float> vector)
        {
            Vector128<float> shuf = Sse.Shuffle(vector, vector, 0x1B); // [3,2,1,0] -> [0,1,2,3]
            Vector128<float> sum1 = Sse.Add(vector, shuf);
            Vector128<float> shuf2 = Sse.Shuffle(sum1, sum1, 0x02); // [1,0,X,X]
            Vector128<float> sum2 = Sse.Add(sum1, shuf2);
            return sum2.ToScalar();
        }

        /// <summary>
        /// 检查SIMD支持情况
        /// </summary>
        public static string GetSIMDSupport()
        {
            if (Avx.IsSupported)
                return "x64 AVX/AVX2";
            else if (AdvSimd.IsSupported)
                return "ARM64 NEON";
            else if (Sse41.IsSupported)
                return "x64 SSE4.1";
            else if (Sse.IsSupported)
                return "x64 SSE";
            else
                return "Scalar (no SIMD)";
        }

        /// <summary>
        /// 获取当前平台的向量处理能力
        /// </summary>
        public static (string Architecture, int VectorSize, string Instructions) GetPlatformInfo()
        {
            if (Avx.IsSupported)
                return ("x64", AVX_VECTOR_SIZE, "AVX/AVX2");
            else if (AdvSimd.IsSupported)
                return ("ARM64", NEON_VECTOR_SIZE, "NEON");
            else if (Sse.IsSupported)
                return ("x64", SSE_VECTOR_SIZE, "SSE");
            else
                return (Environment.Is64BitProcess ? "64-bit" : "32-bit", 1, "Scalar");
        }

        #endregion
    }

   
}
